package com.mcp.client;

import com.mcp.client.transport.*;
import com.mcp.server.model.McpRequest;
import com.mcp.server.model.McpResponse;
import com.mcp.util.JsonUtils;

import java.util.HashMap;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

/**
 * MCP协议客户端
 * 支持多种通信方式
 */
public class McpClient {
    private final McpTransport transport;

    public enum TransportType {
        STDIO, SOCKET, HTTP_SSE
    }

    public McpClient(TransportType type) throws Exception {
        switch (type) {
            case STDIO:
                this.transport = new StdioTransport();
                break;
            // case SOCKET:
            //     this.transport = new SocketTransport();
            //     break;
            case HTTP_SSE:
                this.transport = new HttpSseTransport();
                break;
            default:
                throw new IllegalArgumentException("不支持的传输类型: " + type);
        }
    }

    public McpClient(McpTransport transport) {
        this.transport = transport;
    }

    /**
     * 同步调用远程方法
     */
    public McpResponse call(String method, Object params) throws Exception {
        McpRequest request = new McpRequest();
        request.setMethod(method);
        request.setParams(params);
        request.setId(System.currentTimeMillis()); // 使用时间戳作为ID

        String requestJson = JsonUtils.toJson(request);
        System.out.println("客户端发送请求: " + requestJson);

        String responseJson = transport.send(requestJson);
        System.out.println("客户端接收到响应: " + responseJson);

        return JsonUtils.fromJson(responseJson, McpResponse.class);
    }

    /**
     * 异步调用远程方法（仅适用于HTTP+SSE传输）
     */
    public CompletableFuture<McpResponse> callAsync(String method, Object params) throws Exception {
        if (!(transport instanceof HttpSseTransport)) {
            throw new UnsupportedOperationException("异步调用仅支持HTTP+SSE传输");
        }

        HttpSseTransport httpTransport = (HttpSseTransport) transport;
        CompletableFuture<McpResponse> future = new CompletableFuture<>();

        // 生成唯一请求ID
        String requestId = UUID.randomUUID().toString();

        // 注册响应处理器
        httpTransport.registerResponseHandler(requestId, response -> {
            future.complete(response);
        });

        // 发送请求
        McpRequest request = new McpRequest();
        request.setMethod(method);
        request.setParams(params);
        request.setId(requestId);

        String requestJson = JsonUtils.toJson(request);
        System.out.println("客户端异步发送请求: " + requestJson);

        // 启动请求线程
        new Thread(() -> {
            try {
                transport.send(requestJson);
            } catch (Exception e) {
                future.completeExceptionally(e);
            }
        }).start();

        return future;
    }

    /**
     * 示例用法
     */
    public static void main(String[] args) {
        try {
            // 使用HTTP+SSE传输
            McpClient client = new McpClient(TransportType.HTTP_SSE);

            // 同步调用add方法
            Map<String, Integer> mathParams = new HashMap<>();
            mathParams.put("a", 5);
            mathParams.put("b", 3);
            McpResponse response = client.call("add", mathParams);
            System.out.println("同步调用: 5 + 3 = " + response.getResult());

            // 异步调用add方法
            CompletableFuture<McpResponse> asyncResponse = client.callAsync("add", mathParams);
            System.out.println("异步调用已发送，等待响应...");

            // 等待异步响应，设置超时时间
            try {
                McpResponse asyncResult = asyncResponse.get(5, TimeUnit.SECONDS);
                System.out.println("异步调用结果: 5 + 3 = " + asyncResult.getResult());
            } catch (TimeoutException e) {
                System.err.println("异步调用超时");
            } catch (ExecutionException e) {
                System.err.println("异步调用失败: " + e.getCause().getMessage());
            }

        } catch (Exception e) {
            System.err.println("调用失败: " + e.getMessage());
            e.printStackTrace();  // 打印完整堆栈跟踪
        }
    }
}